
#include "geglu_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling_common.h"
namespace optiling
{
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    gegluTilingData tiling;
    set_size_by_input(tiling, context, 0);
    auto shape = context->GetInputShape(0)->GetStorageShape();

    // 判断dim是否符合规范
    auto finalDim = shape.GetDim(shape.GetDimNum() - 1);
    auto halfFinalDim = finalDim / 2;
    if (finalDim % 2 != 0)
    {
        std::cerr << "finalDim = " << finalDim << " 不能被2整除" << std::endl;
        return ge::GRAPH_FAILED;
    }
    tiling.set_finalDim(finalDim);
    // 获取硬件基本信息
    COMMON_TILING_GET_PLATFORM;
    COMMON_TILING_GET_DTYPE_DATA(0);
    COMMON_TILING_GET_UB_SIZE;
    // 确定计算时需要的ub_buffer如何切分
    auto inputSplitAlignSize = finalDim * dtypeSize;
    auto alignToByteSize = 32;
    auto alignToSize = alignToByteSize / dtypeSize;
    // 分割ub,计算出tileLenth需要的数据
    auto xPieceDatasize = inputSplitAlignSize / 2;
    auto xPieceAlignSize = ALIGN_TO(xPieceDatasize, alignToByteSize);
    auto deltaSize = xPieceAlignSize - xPieceDatasize;
    // 假设finalDim数据量比较小，ub可以容下多个finalDim的数据
    // (x0,x1,y) * double buffer = 3 * 2 = 6
    auto split_ub_size = ub_size / 6;
    // float16类型时不要类型转换，需要弥补12个deltaSize
    // sizeof(float) / sizeof(float16) * 3(x0,x1,y) + 6 = 12
    if (dtype == ge::DataType::DT_FLOAT16)
    {
        split_ub_size = ub_size / 12;
    }
    split_ub_size = split_ub_size / alignToByteSize * alignToByteSize;
    auto splitUbContainDataCount = split_ub_size / dtypeSize;

    // 分割输入数据
    auto blockDim = common_tiling_set_kernel_vec(tiling, context);
    COMMON_TILING_ALIGN_DATA(inputSplitAlignSize, tiling.get_size());
    COMMON_TILING_SPLIT(tiling);

    // 如果ub容不下一份finalDim的数据
    if (xPieceDatasize > split_ub_size)
    {
        tiling.set_largeFinal(true);
        auto tileLength = splitUbContainDataCount;
        tiling.set_formerTileLength(tileLength);
        tiling.set_tailTileLength(tileLength);
        auto finalCalcCount = halfFinalDim % tileLength;
        finalCalcCount = finalCalcCount == 0 ? tileLength : finalCalcCount;
        tiling.set_formerFinalCalcCount(finalCalcCount);
        tiling.set_tailFinalCalcCount(finalCalcCount);
    }
    else
    {
        tiling.set_largeFinal(false);
        auto splitUbHalfFinalCount = splitUbContainDataCount / halfFinalDim;

        auto formerHalfFinalCount = (formerLength / 2) / halfFinalDim;
        auto formerTileLength = formerHalfFinalCount / splitUbHalfFinalCount * halfFinalDim;
        formerTileLength = formerTileLength == 0 ? formerHalfFinalCount * halfFinalDim : formerTileLength;
        tiling.set_formerTileLength(formerTileLength);

        auto tailHalfFinalCount = (tailLength / 2) / halfFinalDim;
        auto tailTileLength = tailHalfFinalCount / splitUbHalfFinalCount * halfFinalDim;
        tailTileLength = tailTileLength == 0 ? tailHalfFinalCount * halfFinalDim : tailTileLength;
        tiling.set_tailTileLength(tailTileLength);

        auto formerFinalCalcCount = (formerHalfFinalCount % splitUbHalfFinalCount) * halfFinalDim;
        formerFinalCalcCount = formerFinalCalcCount == 0 ? formerTileLength : formerFinalCalcCount;
        tiling.set_formerFinalCalcCount(formerFinalCalcCount);

        auto tailFinalCalcCount = (tailHalfFinalCount % splitUbHalfFinalCount) * halfFinalDim;
        tailFinalCalcCount = tailFinalCalcCount == 0 ? tailTileLength : tailFinalCalcCount;
        tiling.set_tailFinalCalcCount(tailFinalCalcCount);
    }

    tiling_zero_workspace(context);
    tiling_save_to_buffer(tiling, context);
    print_common_tiling(tiling);
    std::cout << "finalDim : " << finalDim << std::endl;
    std::cout << "largeFinal : " << tiling.get_largeFinal() << std::endl;
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge
{
static ge::graphStatus InferShape(gert::InferShapeContext *context)
{
    const gert::Shape *x1_shape = context->GetInputShape(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    auto finalIndex = y_shape->GetDimNum() - 1;
    y_shape->SetDim(finalIndex, y_shape->GetDim(finalIndex) / 2);
    return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops
{
class geglu : public OpDef
{
  public:
    explicit geglu(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(geglu);
} // namespace ops
